__author__ = 'sdjkx'


import os
from PIL import Image
from libs.ustr import *

try:
  import xml.etree.cElementTree as ET
except ImportError:
  import xml.etree.ElementTree as ET
from PyQt5.QtWidgets import QMessageBox,QFileDialog

class ScanXmlCp:

    def __init__(self, labelImg, defaultSaveDir):
        self.labelImg = labelImg
        self.defaultSaveDir  = defaultSaveDir
        self.ImageWidth = 0



    def scanCPLabelFun(self):
        if self.defaultSaveDir is not None:
            path = ustr(self.defaultSaveDir)
            ret = QMessageBox.question(self.labelImg, '提示', "确定要扫描目录【"+path+"】里的车牌和公交车线路吗？",
                                       QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
            if ret == QMessageBox.Yes:
                self.scanCpAndBusLabel(path)
                return
        else:
            path = '.'

        dirpath = ustr(QFileDialog.getExistingDirectory(self.labelImg,
                                                       "选择XML文件路径", path,  QFileDialog.ShowDirsOnly
                                                       | QFileDialog.DontResolveSymlinks))

        if dirpath is not None and len(dirpath) > 1:
            ret = QMessageBox.question(self.labelImg, '提示', "确定要扫描目录【"+dirpath+"】里的车牌和公交车线路吗？",
                                       QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
            if ret == QMessageBox.Yes:
                self.scanCpAndBusLabel(dirpath)
                return

    def scanCpAndBusLabel(self, xmlPath):
        basePath =  xmlPath.replace("Annotations","")

        srcDirPath = xmlPath.replace("Annotations","cut")
        if not os.path.exists(srcDirPath):
            srcDirPath = xmlPath.replace("Annotations","src")

        trainDirPath = xmlPath.replace("Annotations","train")
        if not srcDirPath.endswith(os.sep):
            srcDirPath = srcDirPath + "/"

        if not os.path.exists(srcDirPath):
            ret = QMessageBox.question(self.labelImg, '提示', "XML文件对应的原图片目录src不存在，确定要扫描train目录的文件吗？",
                                       QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
            if ret != QMessageBox.Yes:
                return
            else:
                srcDirPath = xmlPath.replace("Annotations","train")

        xmlFileList = os.listdir(xmlPath)
        for xmlFileName in xmlFileList:
            xmlFilePath =  os.path.join(xmlPath, xmlFileName)
            srcFilePath =  os.path.join(srcDirPath , xmlFileName.split(".")[0] + ".jpg")

            if os.path.exists(xmlFilePath) and os.path.exists(srcFilePath):
                if self.ImageWidth == 0:
                    trainImagePath = os.path.join(trainDirPath,xmlFileName.split(".")[0] + ".jpg")
                    trainImg = Image.open(trainImagePath)
                    self.ImageWidth, h = trainImg.size

                tree = ET.parse(xmlFilePath)     # 打开xml文档
                root = tree.getroot()         #获得root节点

                sImg = Image.open(srcFilePath)
                w, h = sImg.size
                scale = w / self.ImageWidth

                for object in root.findall('object'): #找到root节点下的所有object节点
                    label_name = object.find('name').text   # 子节点下节点name的值

                    if label_name == "车牌" or label_name== "公交车路线" or label_name == "红绿灯":
                        bndbox = object.find('bndbox')      # 子节点下属性bndbox的值
                        xmin = float(bndbox.find('xmin').text)
                        ymin = float(bndbox.find('ymin').text)
                        xmax = float(bndbox.find('xmax').text)
                        ymax = float(bndbox.find('ymax').text)

                        left = int(xmin * scale)
                        upper = int(ymin * scale)
                        right = int(xmax * scale)
                        lower = int(ymax * scale)
                        if left < 0:
                            left = 0
                        if upper < 0:
                            upper = 0
                        if right > w:
                            right = w
                        if lower > h:
                            lower = h

                        if left > right:
                            left,right = right,left
                        if upper > lower:
                            upper, lower = lower, upper

                        if label_name == "车牌":
                            destPath = basePath + "Licenselate/"
                        elif label_name == "红绿灯":
                            destPath = basePath + "TrafficSignal/"
                        else:
                            destPath = basePath + "BusRoutes/"

                        if not os.path.exists(destPath):
                            os.makedirs(destPath)

                        destPath = destPath +  xmlFileName.split(".")[0] + '_' + str(left) + ".jpg"
                        if not os.path.exists(destPath):
                            cropped = sImg.crop((left, upper, right, lower))
                            cropped.save(destPath)

